#include "mr_genotyper.h"

const fix_t matching_penalties[80*3] = {
-1, -1, -1,
-45011, -41190, -37821,
-28371, -35263, -44375,
-19795, -31301, -50928,
-14448, -28506, -57482,
-10818, -26466, -64036,
-8232, -24944, -70589,
-6333, -23790, -77143,
-4910, -22906, -83696,
-3829, -22222, -90250,
-2998, -21691, -96804,
-2355, -21276, -103357,
-1854, -20950, -109911,
-1462, -20695, -116464,
-1155, -20493, -123018,
-914, -20334, -129572,
-723, -20208, -136125,
-573, -20108, -142679,
-454, -20030, -149232,
-360, -19967, -155786,
-285, -19918, -162340,
-226, -19878, -168893,
-179, -19847, -175447,
-142, -19823, -182000,
-113, -19803, -188554,
-89, -19787, -195108,
-71, -19775, -201661,
-56, -19765, -208215,
-44, -19757, -214768,
-35, -19751, -221322,
-27, -19746, -227876,
-22, -19742, -234429,
-17, -19739, -240983,
-13, -19737, -247536,
-10, -19735, -254090,
-8, -19733, -260644,
-6, -19732, -267197,
-5, -19731, -273751,
-4, -19730, -280304,
-3, -19730, -286858,
-2, -19729, -293412,
-1, -19729, -299965,
-1, -19728, -306519,
0, -19728, -313072,
0, -19728, -319626,
0, -19728, -326180,
0, -19728, -332733,
0, -19728, -339287,
0, -19728, -345840,
0, -19728, -352394,
0, -19727, -358948,
0, -19727, -365501,
0, -19727, -372055,
0, -19727, -378608,
0, -19727, -385162,
0, -19727, -391716,
0, -19727, -398269,
0, -19727, -404823,
0, -19727, -411376,
0, -19727, -417930,
0, -19727, -424484,
0, -19727, -431037,
0, -19727, -437591,
0, -19727, -444144,
0, -19727, -450698,
0, -19727, -457252,
0, -19727, -463805,
0, -19727, -470359,
0, -19727, -476912,
0, -19727, -483466,
0, -19727, -490020,
0, -19727, -496573,
0, -19727, -503127,
0, -19727, -509680,
0, -19727, -516234,
0, -19727, -522788,
0, -19727, -529341,
0, -19727, -535895,
0, -19727, -542448,
0, -19727, -549002
};

//Assumes phred is a multiple of 10
fix_t q_base_given_genotype(phred_score_t phred, int mismatchA, int mismatchB){
  int i, j;
  fix_t toRet = 0;
  for(i = 10; i <= 70; i+=10){
    fix_t penalty = 0;
    if (i == phred){
      //0 mismatches - full, 1 mismatch - half, 
      // 2 mismatches - none
      if (mismatchA){
        if (mismatchB){
          penalty = matching_penalties[i*3 + 2];
        } else {
          penalty = matching_penalties[i*3 + 1];
        }
      } else {
        if (mismatchB){
          penalty = matching_penalties[i*3 + 1];
        } else {
          penalty = matching_penalties[i*3 + 0];
        }
      }
    }
    toRet += penalty;
  }
  return toRet;
}

const fix_t reference_priors_lut[4][NUM_GENOTYPES] = {
  //Generated by the call
  //DiploidGenotypePriors.getReferencePolarizedPriors(ref,
  //            DiploidGenotypePriors.HUMAN_HETEROZYGOSITY,
  //                        0.01);
  //in the GATK and rounding the results to fix_t
  //A
  {-42, -228162, -228162, -228162, -247604, -358948, -358948, -247604,
  -358948, -247604, },
  //T
  {-247604, -358948, -358948, -228162, -247604, -358948, -228162,
  -247604, -228162, -42, },
  //G
  {-247604, -358948, -228162, -358948, -247604, -228162, -358948, -42,
  -228162, -247604, },
  //C
  {-247604, -228162, -358948, -358948, -42, -228162, -228162, -247604,
  -358948, -247604, },
};

void getReferencePolarizedPriors(fix_t* priors, nucleotide_base_t ref){
  nucleotide_base_t i;
  int j;
  for(i = 0; i < 4; i++){
    if (i == ref){
      for(j = 0; j < NUM_GENOTYPES; j++){
        priors[j] = reference_priors_lut[i][j];
      }
    }
  }
}

typedef struct _top_two_calls {
  single_locus_genotype_t best_call;
  fix_t qual_best_call;
  fix_t qual_second_best_call;
} top_two_calls;

void consider_possibility(top_two_calls* caller, single_locus_genotype_t call, fix_t likelihood){
  //Is this more likely than the current best?
  if (likelihood > caller->qual_best_call){
    //Shift both down
    caller->best_call = call;
    caller->qual_second_best_call = caller->qual_best_call;
    caller->qual_best_call = likelihood;
  } else {
    //Is this more likely than the second best?
    if (likelihood > caller->qual_second_best_call){
      caller->qual_second_best_call = likelihood;
    }
  }
}

void map (MapperIn *mapper_input, MapperOut *mapper_output) {
  int i;

  MapperChunkOut* out = &mapper_output->output[0];

  //Only one reducer.
  assert_zero(NUM_REDUCERS-1);

  //Iterate over pileups, making genotyping calls
  for(i = 0; i < NUM_LOCI_PER_MAPPER; i++){
    //Compute likelihoods that the organism sequenced has a particular
    //genotype at this locus
    fix_t likelihood [NUM_GENOTYPES];

    //Keep track of the top two possibilities
    top_two_calls caller = {0, FIX_T_MIN, FIX_T_MIN};

    //Information: The reference base for this locus
    nucleotide_base_t ref = mapper_input->refseq[i];
    getReferencePolarizedPriors(likelihood, ref);

    single_locus_genotype_t j;
    for(j = 0; j < NUM_GENOTYPES; j++){
      //More information: the sequence reads at this locus
      int u;
      for (u = 0; u < MAX_SIZE_PILEUP; u++){
        phred_score_t qual = mapper_input->data[i].pileup[u].quality;
        //Truncate qual to the next lowest multiple of 10
        qual -= qual % 10; //Optimization.

        nucleotide_base_t base = mapper_input->data[i].pileup[u].base;

        //If qual is too small (for now, ignore 0's), ignore the base.
        if (qual != 0){
          likelihood[j] += q_base_given_genotype(qual,
            base != GENOTYPES[j][0],
            base != GENOTYPES[j][1]);
        }
      }

      //Consider this possibility.
      consider_possibility(&caller, j, likelihood[j]);
    }

    //Make the call based on the top two likelihoods
    single_locus_genotype_call* call = out->calls + i;
    call->call = caller.best_call;
    //Call quality is the LOD between the best call and second call
    call->quality = caller.qual_best_call - caller.qual_second_best_call;
  }
}

// include this header _only_ after the above are defined
#include <mapred_map.h>
